import numpy
from projects.electrophysiology.view_data import frame_utils
        
def sta(frames, frame_times, spike_times, stim_onset_time, tau):
    """
    Calculate the spike triggered average.
    Inputs:
        frames           : A 3D numpy array of formated frames. 
                           (see format_frames)
        frame_times      : A list of times when the frames were drawn.
        spike_times      : A list of times when 'events' or spikes occured
        tau              : the window of time BEFORE an 'event' to look at 
                           the stimulus.
        stim_onset_time  : the timing of the stimulus
    Returns:
        sta              : A 2D(greyscale) or 3D(color) numpy array of 
                           formated frames 
                           depending on the format of the frames that were
                           inputs.
    """
    
    spt = numpy.array(spike_times) - stim_onset_time - tau 
    #shifted the spike times to stimulus onset and the delay tau
    spt = spt[spt<frame_times[-1]] #get rid of spikes after frames stop 
    shifted_spike_times = spt[spt>=0] #get rid of spikes before frames began
    
    spike_trig_avg = numpy.zeros(frames[0].shape, dtype=numpy.float64)
    for time in shifted_spike_times:
        spike_trig_avg += frame_utils.get_stim(time, frames, frame_times, 
                                               None, None)
    spike_trig_avg /= len(shifted_spike_times)
    return spike_trig_avg, len(shifted_spike_times)



